#!/usr/bin/env python3
#-*- coding:utf8 -*-
# Power by 2020-06-03 23:30:40

import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Linear
import numpy as np
import os
import json
import random
import gzip

def load_data(model="train"):
    """TODO: Docstring for load_data.
    :returns: TODO

    """
    data_file='./mnist.json.gz' 
    data=json.load(gzip.open(data_file))
    train_set,val_set,test_set=data
    print('mnist data load done')
    if model == 'train':
        imgs,labels=train_set[0],train_set[1]
    elif model == 'valid':
        imgs,labels=val_set[0],val_set[1]
    elif model == 'eval':
        imgs,labels=test_set[0],test_set[1]
    else:
        raise Exception("model can only be one of ['train''valid''eval']")
    print('训练数据集数量:',len(imgs))
    imgs_length=len(imgs)
    assert len(imgs)==len(labels),"length of train_imgs({}) should be the same as train_labels({})".format(len(imgs),len(labels))
    imgs_length=len(imgs)
    index_list = list(range(imgs_length))
    BATCHSIZE=100
    IMG_ROWS = 28
    IMG_COLS = 28

    def data_generator():
        if model=='train':
            random.shuffle(index_list)
        imgs_list=[]
        labels_list=[]
        for i in index_list:
            img=np.reshape(imgs[i],[1,IMG_ROWS,IMG_COLS]).astype('float32')
            label=np.reshape(labels[i],[1]).astype('float32')
            imgs_list.append(img)
            labels_list.append(label)
            if len(imgs_list)==BATCHSIZE:
                yield np.array(imgs_list),np.array(labels_list)
                imgs_list=[]
                labels_list=[]
        if len(imgs_list) > 0:
            yield np.array(imgs_list),np.array(labels_list)
    return data_generator

